///////////////////////////////////////////////////////////////////////////////
//
// Module Name:
//
//     dump.c
//
// Abstract:
//
//     Main code that is run when the user clicks "Generate"
//
// Author:
//
//     vxiiduu (29-Oct-2022)
//
// Revision History:
//
//     vxiiduu               29-Oct-2022  Initial creation.
//
///////////////////////////////////////////////////////////////////////////////

#include "buildcfg.h"
#include <KexComm.h>
#include "resource.h"
#include "kexexprt.h"

LARGE_UNICODE_STRING Output;

STATIC VOID AppendToOutput(
	IN	PCWSTR	Format,
	IN	...)
{
	HRESULT Result;
	ARGLIST ArgList;
	SIZE_T TextCch;
	PWSTR Text;

	va_start(ArgList, Format);

	Result = StringCchVPrintfBufferLength(&TextCch, Format, ArgList);
	if (FAILED(Result)) {
		return;
	}

	Text = StackAlloc(WCHAR, TextCch);
	Result = StringCchVPrintf(Text, TextCch, Format, ArgList);
	
	if (FAILED(Result)) {
		return;
	}

	if (!Output.Buffer) {
		Output.MaxCch = 1024;
		Output.Buffer = SafeAlloc(WCHAR, Output.MaxCch);
		Output.Cch = 0;
	}

	if (Output.Cch + TextCch > Output.MaxCch) {
		Output.MaxCch *= 2;
		Output.Buffer = SafeReAlloc(Output.Buffer, WCHAR, Output.MaxCch);
	}

	if (!Output.Buffer) {
		return;
	}

	RtlCopyMemory(Output.Buffer + Output.Cch, Text, TextCch * sizeof(WCHAR));
	Output.Cch += (ULONG) TextCch - 1;
}

VOID DumpExports(
	IN	HWND					MainWindow,
	IN	PCWSTR					DllPathCString,
	IN	KEXEXPRT_GENERATE_STYLE	Style,
	IN	BOOLEAN					IncludeOrdinals)
{
	BOOLEAN Success;
	NTSTATUS Status;

	UNICODE_STRING FilePath;
	OBJECT_ATTRIBUTES ObjectAttributes;
	IO_STATUS_BLOCK IoStatusBlock;
	HANDLE FileHandle;
	HANDLE SectionHandle;

	UNICODE_STRING DllBaseName;
	UNICODE_STRING DotDll;
	PVOID DllBase;
	SIZE_T DllSize;
	PIMAGE_EXPORT_DIRECTORY ExportDirectory;
	ULONG ImportDescriptorSize;
	PDWORD NameRvas;
	PUSHORT NameOrdinals;
	PBOOLEAN OrdinalHasName;
	ULONG Index;

	DllBase = NULL;
	DllSize = 0;
	FileHandle = NULL;
	SectionHandle = NULL;

	Success = RtlDosPathNameToRelativeNtPathName_U(
		DllPathCString,
		&FilePath,
		NULL,
		NULL);

	if (!Success) {
		ErrorBoxF(L"Failed to convert the input file name from DOS to NT.");
		return;
	}

	InitializeObjectAttributes(&ObjectAttributes, &FilePath, OBJ_CASE_INSENSITIVE, NULL, NULL);

	Status = NtOpenFile(
		&FileHandle,
		FILE_GENERIC_READ,
		&ObjectAttributes,
		&IoStatusBlock,
		FILE_SHARE_READ,
		FILE_SYNCHRONOUS_IO_NONALERT);
	
	if (!NT_SUCCESS(Status)) {
		ErrorBoxF(L"Failed to open the specified DLL file \"%wZ\": 0x%08lx: %s",
				  &FilePath, Status, KexRtlNtStatusToString(Status));
		goto Exit;
	}

	Status = NtCreateSection(
		&SectionHandle,
		SECTION_MAP_READ,
		NULL,
		0,
		PAGE_READONLY,
		SEC_IMAGE,
		FileHandle);

	if (!NT_SUCCESS(Status)) {
		ErrorBoxF(L"Failed to create section: 0x%08lx: %s",
				  Status, KexRtlNtStatusToString(Status));
		goto Exit;
	}

	Status = NtMapViewOfSection(
		SectionHandle,
		NtCurrentProcess(),
		&DllBase,
		0,
		0,
		NULL,
		&DllSize,
		ViewUnmap,
		0,
		PAGE_READONLY);

	if (!NT_SUCCESS(Status)) {
		ErrorBoxF(L"Failed to map DLL: 0x%08lx: %s",
				  Status, KexRtlNtStatusToString(Status));
		goto Exit;
	}

	ExportDirectory = (PIMAGE_EXPORT_DIRECTORY) RtlImageDirectoryEntryToData(
		DllBase,
		TRUE,
		IMAGE_DIRECTORY_ENTRY_EXPORT,
		&ImportDescriptorSize);

	KexRtlPathFindFileName(&FilePath, &DllBaseName);
	NameRvas = (PDWORD) RVA_TO_VA(DllBase, ExportDirectory->AddressOfNames);
	NameOrdinals = (PUSHORT) RVA_TO_VA(DllBase, ExportDirectory->AddressOfNameOrdinals);
	OrdinalHasName = StackAlloc(BOOLEAN, ExportDirectory->NumberOfFunctions);
	RtlZeroMemory(OrdinalHasName, ExportDirectory->NumberOfFunctions * sizeof(BOOLEAN));

	RtlZeroMemory(&Output, sizeof(Output));

	switch (Style) {
	case GenerateStyleDef:
		AppendToOutput(L";; Generated by KexExprt from \"%s\"\r\n", DllPathCString);
		AppendToOutput(L"LIBRARY %wZ\r\n", &DllBaseName);
		AppendToOutput(L"EXPORTS\r\n");
		break;
	case GenerateStylePragma:
		AppendToOutput(L"// Generated by KexExprt from \"%s\"\r\n", DllPathCString);
		break;
	default:
		NOT_REACHED;
	}

	RtlInitConstantUnicodeString(&DotDll, L".dll");
	
	if (KexRtlUnicodeStringEndsWith(&DllBaseName, &DotDll, TRUE)) {
		DllBaseName.Length -= 4 * sizeof(WCHAR);
		DllBaseName.Buffer[DllBaseName.Length / sizeof(WCHAR)] = '\0';
	}
	
	for (Index = 0; Index < ExportDirectory->NumberOfNames; ++Index) {
		PCSTR FunctionName;
		ULONG Ordinal;

		Ordinal = NameOrdinals[Index] + ExportDirectory->Base;
		FunctionName = (PCSTR) RVA_TO_VA(DllBase, NameRvas[Index]);
		OrdinalHasName[NameOrdinals[Index]] = TRUE;

		switch (Style) {
		case GenerateStyleDef:
			if (IncludeOrdinals) {
				AppendToOutput(
					L"\t%hs = %wZ.%hs @%lu\r\n",
					FunctionName, &DllBaseName, FunctionName, Ordinal);
			} else {
				AppendToOutput(
					L"\t%hs = %wZ.%hs\r\n",
					FunctionName, &DllBaseName, FunctionName);
			}
			break;
		case GenerateStylePragma:
			if (IncludeOrdinals) {
				AppendToOutput(
					L"#pragma comment(linker, \"/EXPORT:%hs=%wZ.%hs,@%lu\")\r\n",
					FunctionName, &DllBaseName, FunctionName, Ordinal);
			} else {
				AppendToOutput(
					L"#pragma comment(linker, \"/EXPORT:%hs=%wZ.%hs\")\r\n",
					FunctionName, &DllBaseName, FunctionName);
			}
			break;
		default:
			NOT_REACHED;
		}
	}

	for (Index = 0; Index < ExportDirectory->NumberOfFunctions; ++Index) {
		ULONG Ordinal;

		Ordinal = Index + ExportDirectory->Base;

		if (!OrdinalHasName[Index]) {
			switch (Style) {
			case GenerateStyleDef:
				AppendToOutput(L"\t__OrdinalFunction%lu = %wZ.#%lu @%lu NONAME\r\n", Ordinal, &DllBaseName, Ordinal, Ordinal);
				break;
			case GenerateStylePragma:
				AppendToOutput(L"#pragma comment(linker, \"/EXPORT:__OrdinalFunction%lu=%wZ.#%lu,@%lu,NONAME\")\r\n", Ordinal, &DllBaseName, Ordinal, Ordinal);
				break;
			default:
				NOT_REACHED;
			}
		}
	}
	
	SetDlgItemText(MainWindow, IDC_RESULT, Output.Buffer);

Exit:
	SafeFree(Output.Buffer);
	SafeFree(FilePath.Buffer);

	if (DllBase) {
		NtUnmapViewOfSection(NtCurrentProcess(), DllBase);
	}

	SafeClose(FileHandle);
	SafeClose(SectionHandle);
}